// Usage sudo ./rx local_ip <device> (device is optional)
// e.g. sudo ./rx 192.168.1.2

#include <stdio.h>
#include <stdlib.h>
#include <assert.h>
#include <string.h>
#include <ctype.h>
#include <time.h>
#include <pcap.h>

#define __FAVOR_BSD

#include <netinet/in.h>
#include <netinet/ip.h>
#include <netinet/tcp.h>
#include <netinet/udp.h>
#include <arpa/inet.h>

/* default snap length (maximum bytes per packet to capture) */
#define SNAP_LEN 1518

/* ethernet headers are always exactly 14 bytes [1] */
#define SIZE_ETHERNET 14

/* Ethernet addresses are 6 bytes */
#define ETHER_ADDR_LEN  6

/* Ethernet header */
struct sniff_ethernet {
  u_char  ether_dhost[ETHER_ADDR_LEN];    /* destination host address */
  u_char  ether_shost[ETHER_ADDR_LEN];    /* source host address */
  u_short ether_type;                     /* IP? ARP? RARP? etc */
};

/* IP header */
struct sniff_ip {
  u_char  ip_vhl;                 /* version << 4 | header length >> 2 */
  u_char  ip_tos;                 /* type of service */
  u_short ip_len;                 /* total length */
  u_short ip_id;                  /* identification */
  u_short ip_off;                 /* fragment offset field */
  #define IP_RF 0x8000            /* reserved fragment flag */
  #define IP_DF 0x4000            /* dont fragment flag */
  #define IP_MF 0x2000            /* more fragments flag */
  #define IP_OFFMASK 0x1fff       /* mask for fragmenting bits */
  u_char  ip_ttl;                 /* time to live */
  u_char  ip_p;                   /* protocol */
  u_short ip_sum;                 /* checksum */
  struct  in_addr ip_src,ip_dst;  /* source and dest address */
};
#define IP_HL(ip)               (((ip)->ip_vhl) & 0x0f)
#define IP_V(ip)                (((ip)->ip_vhl) >> 4)

/* TCP header */
typedef u_int tcp_seq;

struct sniff_tcp {
  u_short th_sport;               /* source port */
  u_short th_dport;               /* destination port */
  tcp_seq th_seq;                 /* sequence number */
  tcp_seq th_ack;                 /* acknowledgement number */
  u_char  th_offx2;               /* data offset, rsvd */
#define TH_OFF(th)      (((th)->th_offx2 & 0xf0) >> 4)
  u_char  th_flags;
  #define TH_FIN  0x01
  #define TH_SYN  0x02
  #define TH_RST  0x04
  #define TH_PUSH 0x08
  #define TH_ACK  0x10
  #define TH_URG  0x20
  #define TH_ECE  0x40
  #define TH_CWR  0x80
  #define TH_FLAGS        (TH_FIN|TH_SYN|TH_RST|TH_ACK|TH_URG|TH_ECE|TH_CWR)
  u_short th_win;                 /* window */
  u_short th_sum;                 /* checksum */
  u_short th_urp;                 /* urgent pointer */
};

long ns_offset;
long long get_time_ns() {
  struct timespec ts;
  clock_gettime(CLOCK_REALTIME, &ts);
  long long time = ((long long)ts.tv_sec) * 1000000000 + ts.tv_nsec;
  return time;
}

void print_hex_ascii_line(const u_char *payload, int len, int offset)
{
  int i;
  int gap;
  const u_char *ch;

  /* offset */
  printf("%05d   ", offset);
  
  /* hex */
  ch = payload;
  for(i = 0; i < len; i++) {
    printf("%02x ", *ch);
    ch++;
    /* print extra space after 8th byte for visual aid */
    if (i == 7)
      printf(" ");
  }
  /* print space to handle line less than 8 bytes */
  if (len < 8)
    printf(" ");
  
  /* fill hex gap with spaces if not full line */
  if (len < 16) {
    gap = 16 - len;
    for (i = 0; i < gap; i++) {
      printf("   ");
    }
  }
  printf("   ");
  
  /* ascii (if printable) */
  ch = payload;
  for(i = 0; i < len; i++) {
    if (isprint(*ch))
      printf("%c", *ch);
    else
      printf(".");
    ch++;
  }

  printf("\n");
}

void print_payload(const u_char *payload, int len)
{
  u_int64_t time = -1;
  if (len <= 0) {
    printf("FAIL,");
    return;
  }
  
  memcpy(&time, payload, sizeof(time));
  // timestamp
  printf("%lld,", (long long int)time);
  // true source hostname
  printf("%s,", (char*)((void*)payload+sizeof(time)));
}

void print_usage() {
  fprintf(stderr, "Usage: rx local_ip time_offset <device>\n");
}

char* get_default_device() {
  char *dev;
  char errbuf[PCAP_ERRBUF_SIZE];
  dev = pcap_lookupdev(errbuf);
  if (dev == NULL) {
    fprintf(stderr, "Couldn't find default device: %s\nAre you root?\n", errbuf);
    return NULL;
  }
  return dev;
}

char* get_device(int argc, char** argv) {
  char* dev;
  if (argc > 3) {
    dev = argv[3];
  }
  else {
    dev = get_default_device();
  }
  if (dev == NULL) {
    fprintf(stderr, "No device selected. Quitting.\n");
    print_usage();
    exit(EXIT_FAILURE);
  }
  return dev;
}

char* get_local_ip(int argc, char** argv) {
  char* ip_str = NULL;
  if (argc > 1) {
    ip_str = argv[1];
  }
  if (ip_str == NULL) {
    fprintf(stderr, "No local IP specified. Quitting.\n");
    print_usage();
    exit(EXIT_FAILURE);
  }
  // 255.255.255.255\0
  char* filtered_ip_str = malloc(16*sizeof(char));
  strncpy(filtered_ip_str, ip_str, 15);
  filtered_ip_str[15] = 0;
  return filtered_ip_str;
}

void got_packet(u_char *args, const struct pcap_pkthdr *header, const u_char *packet)
{

  static int count = 1;                   /* packet counter */
  
  /* declare pointers to packet headers */
  //const struct sniff_ethernet *ethernet;  /* The ethernet header [1] */
  const struct sniff_ip *ip;              /* The IP header */
  const struct tcphdr *tcp;            /* The TCP header */
  const struct udphdr *udp;            /* The TCP header */
  const u_char *payload;                    /* Packet payload */
  long long time_received;
  int size_ip;
  int size_tcp;
  int size_payload;
  
  time_received = get_time_ns() + ns_offset;

  //printf("\nPacket number %d:\n", count);
  count++;
  
  /* define ethernet header */
  //ethernet = (struct sniff_ethernet*)(packet);
  
  /* define/compute ip header offset */
  ip = (struct sniff_ip*)(packet + SIZE_ETHERNET);
  size_ip = IP_HL(ip)*4;
  if (size_ip < 20) {
    fprintf(stderr, "   * Invalid IP header length: %u bytes\n", size_ip);
    return;
  }

  /* print source and destination IP addresses */
  // CSV format: ip_src,ip_dst,PROTO,port_src,port_dst,time_sent,hostname,time_received
  //printf("       From: %s\n", inet_ntoa(ip->ip_src));
  //printf("         To: %s\n", inet_ntoa(ip->ip_dst));
  
  printf("%s,%s,", inet_ntoa(ip->ip_src), inet_ntoa(ip->ip_dst));

  /* determine protocol */  
  switch(ip->ip_p) {
    case IPPROTO_TCP:
      printf("TCP,");
      
      tcp = (struct tcphdr*)(packet + SIZE_ETHERNET + size_ip);
      size_tcp = tcp->th_off*4;
      if (size_tcp < 20) {
        printf("   * Invalid TCP header length: %u bytes\n", size_tcp);
        return;
      }
      printf("%d,%d,", ntohs(tcp->th_sport), ntohs(tcp->th_dport));
      payload = (u_char *)(packet + SIZE_ETHERNET + size_ip + size_tcp);
      size_payload = ntohs(ip->ip_len) - (size_ip + size_tcp);
      
      if (size_payload > 0) {
        print_payload(payload, size_payload);
      }
      break;
    case IPPROTO_UDP:
      printf("UDP,");
      
      udp = (struct udphdr*)(packet + SIZE_ETHERNET + size_ip);
      printf("%d,%d,", ntohs(udp->uh_sport), ntohs(udp->uh_dport));
      payload = (u_char*)(packet + SIZE_ETHERNET + size_ip + sizeof(struct udphdr));
      size_payload = ntohs(ip->ip_len) - (size_ip + sizeof(struct udphdr));

      if (size_payload > 0) {
        print_payload(payload, size_payload);
      }
      break;
    case IPPROTO_ICMP:
      printf("ICMP,");
      break;
    case IPPROTO_IP:
      printf("IP,");
      break;
    default:
      printf("UNKNOWN");
      break;
  }
  
  printf("%lld,\n", time_received);
}

void init_pcap_capture(char* dev, char* local_ip) {
  char errbuf[PCAP_ERRBUF_SIZE];
  pcap_t* handle;
  char filter_expression[32];
  bpf_u_int32 net;
  bpf_u_int32 mask;
  struct bpf_program filter_program;
  sprintf(filter_expression, "dst host %s and ip and not (src port 22 or dst port 22)", local_ip);
  if ((handle = pcap_open_live(dev, 1514, 1, 500, errbuf)) == NULL) {
    fprintf(stderr, "Couldn't open device %s: %s\n", dev, errbuf);
    exit(EXIT_FAILURE);
  }
  if (pcap_lookupnet(dev, &net, &mask, errbuf) == -1) {
    fprintf(stderr, "Couldn't get netmask for device %s: %s\n", dev, errbuf);
    exit(EXIT_FAILURE);
  }
  if (pcap_compile(handle, &filter_program, filter_expression, 0, net) == -1) {
    fprintf(stderr, "Couldn't parse filter %s: %s\n", filter_expression, pcap_geterr(handle));
    exit(EXIT_FAILURE);
  }
  if (pcap_setfilter(handle, &filter_program) == -1) {
    fprintf(stderr, "Couldn't install filter %s: %s\n", filter_expression, pcap_geterr(handle));
    exit(EXIT_FAILURE);
  }
  
  printf("Starting capture...\n");
  
  pcap_loop(handle, -1, got_packet, NULL);
  
  // Cleanup
  pcap_freecode(&filter_program);
  pcap_close(handle);
}

int main(int argc, char** argv) {
  char* dev = get_device(argc, argv);
  if (argc < 3) {
    print_usage();
    exit(EXIT_FAILURE);
  }
  float time_offset;
  sscanf(argv[2], "%f", &time_offset);
  ns_offset = time_offset * 1000000000;
  char* local_ip = get_local_ip(argc, argv);
  //printf("The device is %s\n", dev);
  init_pcap_capture(dev, local_ip);
  return 0;
}
